"""
MCP (Model Context Protocol) data models for JSON-RPC 2.0 format
"""
from typing import Optional, Any, Dict, Union, List
from pydantic import BaseModel, Field
from enum import Enum


class MCPMethod(str, Enum):
    """MCP standard methods"""
    TOOLS_LIST = "tools/list"
    TOOLS_CALL = "tools/call"
    INITIALIZE = "initialize"
    INITIALIZED = "initialized"


class MCPError(BaseModel):
    """JSON-RPC 2.0 error object"""
    code: int = Field(..., description="Error code")
    message: str = Field(..., description="Error message")
    data: Optional[Any] = Field(None, description="Additional error data")


class MCPRequest(BaseModel):
    """JSON-RPC 2.0 request format for MCP protocol"""
    jsonrpc: str = Field("2.0", description="JSON-RPC version")
    id: Optional[Union[str, int]] = Field(None, description="Request ID")
    method: MCPMethod = Field(..., description="MCP method name")
    params: Optional[Dict[str, Any]] = Field(None, description="Method parameters")

    class Config:
        use_enum_values = True


class MCPResponse(BaseModel):
    """JSON-RPC 2.0 response format for MCP protocol"""
    jsonrpc: str = Field("2.0", description="JSON-RPC version")
    id: Optional[Union[str, int]] = Field(None, description="Request ID")
    result: Optional[Any] = Field(None, description="Success result")
    error: Optional[MCPError] = Field(None, description="Error object")

    def model_validate(cls, values):
        """Validate that either result or error is present, but not both"""
        if values.get('result') is not None and values.get('error') is not None:
            raise ValueError("Response cannot have both result and error")
        if values.get('result') is None and values.get('error') is None:
            raise ValueError("Response must have either result or error")
        return values


class ToolDefinition(BaseModel):
    """Tool definition for tools/list response"""
    name: str = Field(..., description="Tool name")
    description: str = Field(..., description="Tool description")
    input_schema: Dict[str, Any] = Field(..., description="Tool input JSON schema")


class ToolsListResult(BaseModel):
    """Result format for tools/list method"""
    tools: List[ToolDefinition] = Field(default_factory=list, description="Available tools")


class ToolCallParams(BaseModel):
    """Parameters for tools/call method"""
    name: str = Field(..., description="Tool name to call")
    arguments: Dict[str, Any] = Field(default_factory=dict, description="Tool arguments")


class ToolCallResult(BaseModel):
    """Result format for tools/call method"""
    content: List[Dict[str, Any]] = Field(default_factory=list, description="Tool execution result")
    isError: Optional[bool] = Field(False, description="Whether the tool execution resulted in an error")


class InitializeParams(BaseModel):
    """Parameters for initialize method"""
    protocolVersion: str = Field(..., description="MCP protocol version")
    capabilities: Dict[str, Any] = Field(default_factory=dict, description="Client capabilities")
    clientInfo: Dict[str, str] = Field(default_factory=dict, description="Client information")


class ServerCapabilities(BaseModel):
    """Server capabilities for initialize response"""
    logging: Optional[Dict[str, Any]] = Field(None, description="Logging capabilities")
    prompts: Optional[Dict[str, Any]] = Field(None, description="Prompts capabilities")
    resources: Optional[Dict[str, Any]] = Field(None, description="Resources capabilities")
    tools: Optional[Dict[str, Any]] = Field(None, description="Tools capabilities")


class InitializeResult(BaseModel):
    """Result format for initialize method"""
    protocolVersion: str = Field("2024-11-05", description="MCP protocol version")
    capabilities: ServerCapabilities = Field(default_factory=ServerCapabilities, description="Server capabilities")
    serverInfo: Dict[str, str] = Field(default_factory=dict, description="Server information")


class MCPNotification(BaseModel):
    """JSON-RPC 2.0 notification format (no response expected)"""
    jsonrpc: str = Field("2.0", description="JSON-RPC version")
    method: str = Field(..., description="Notification method")
    params: Optional[Dict[str, Any]] = Field(None, description="Notification parameters")